# coding=utf-8
import sqlite3

import pandas


class MyDB(object):
    def __init__(self, db_name='mystocks.db'):
        self.conn = sqlite3.connect(db_name)
        self.cur = self.conn.cursor()

    def close(self):
        self.cur.close()
        self.conn.close()

    def change_csv2sqlite(self, csv_name, table_name, from_col=0):
        df = pandas.read_csv(csv_name)
        df = df.iloc[:, from_col:]
        df.to_sql(table_name, self.conn, if_exists='replace')

    def select_all_indus_from_db(self):
        result = self.cur.execute('SELECT industry,number FROM industry_count_rank')
        return result

    def select_indus_from_db(self, stock):
        message = self.cur.execute('select * from stocks where ts_code="{}"'.format(stock)).fetchall()
        for i in message:
            return i[5]

    def select_indus_from_db_by_name(self, name):
        message = self.cur.execute('select * from stocks where name="{}"'.format(name)).fetchall()
        for i in message:
            return i[5]

    def select_stocks_from_db_by_industry(self, indus):
        message = self.cur.execute('select * from stocks where industry="{}"'.format(indus)).fetchall()
        stocks = [i[1] for i in message]
        return stocks

    def select_all_stocks_by_name(self, names):
        if len(names) > 1:
            message = self.cur.execute(f'select ts_code from stocks where name in {tuple(names)}').fetchall()
        else:
            message = self.cur.execute(f'select ts_code from stocks').fetchall()
        res = [i[0] for i in message]
        return res

    def select_all_stocks_and_name_by_name(self, names):
        message = self.cur.execute(f'select ts_code,name from stocks where name in {tuple(names)}').fetchall()
        return message

    def select_all_indus_from_db(self):
        message = self.cur.execute('select * from stocks').fetchall()
        res = [i[5] for i in message]
        return res

    def insert2db(self):
        self.cur.execute('INSERT INTO CONFIG VALUES(?,?,?)",()')
        self.conn.commit()
